import torch
import numpy as np
from torch import nn
import torch.nn.functional as F
from typing import Any, Dict, List, Type, Optional, Callable, Union

from tianshou.policy import A2CPolicy
from tianshou.data import Batch, ReplayBuffer, to_torch_as, to_numpy

from ril.policy.rule_policy_base import RulePolicyBase


class A2CRulePolicy(A2CPolicy, RulePolicyBase):
    def __init__(
        self,
        rule_fn: Callable = None,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        A2CPolicy.__init__(self, *args, **kwargs)
        RulePolicyBase.__init__(self, rule_fn)

    def forward(
        self,
        batch: Batch,
        state: Optional[Union[dict, Batch, np.ndarray]] = None,
        **kwargs: Any,
    ) -> Batch:
        obs_, info_ = batch.obs, batch.info
        logits, h = self.actor(obs_, state=state, info=info_)

        if isinstance(logits, tuple):
            dist = self.dist_fn(*logits)
        else:
            dist = self.dist_fn(logits)
        act = dist.sample()

        # use ril
        if not self.updating:
            act, mask, explorable = self.process_rule(batch, act)
            if not np.any(mask == None):
                # this mask come from ril
                if isinstance(logits, tuple):
                    raise Exception("NOT IMPL")
                else:
                    logits_array = to_numpy(logits)
                    logits_array[~mask] = 0
                    if not logits_array.any():
                        logits_array[mask] = 0.1
                    dist = self.dist_fn(to_torch_as(logits_array, logits))
                act = dist.sample()

        return Batch(logits=logits, act=act, state=h, dist=dist)
